#!/usr/bin/env python
# LeCrunch
# Copyright (C) 2010 Anthony LaTorre 
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import sys
import ROOT
ROOT.gROOT.SetStyle("Plain")
import numpy
import rootfile
import subprocess

def persist(file, nevents):
    data = rootfile.RootFile(file)

    mg1 = ROOT.TMultiGraph('mg1', '')
    mg2 = ROOT.TMultiGraph('mg2', '')
    mg3 = ROOT.TMultiGraph('mg3', '')
    mg4 = ROOT.TMultiGraph('mg4', '')

    for n in range(len(data)):
        print '\rScanning event: %i' % (n+1),
        sys.stdout.flush()

        data.getentry(n)
        channels = data.getchannels()
        if 1 in channels:
            x1, y1 = data.getc1()
        else:
            x1, y1 = numpy.array([0]), numpy.array([0])
        if 2 in channels:
            x2, y2 = data.getc2()
        else:
            x2, y2 = numpy.array([0]), numpy.array([0])
        if 3 in channels:
            x3, y3 = data.getc3()
        else:
            x3, y3 = numpy.array([0]), numpy.array([0])
        if 4 in channels:
            x4, y4 = data.getc4()
        else:
            x4, y4 = numpy.array([0]), numpy.array([0])
        
        g1 = ROOT.TGraph(len(x1), x1, y1)
        g2 = ROOT.TGraph(len(x2), x2, y2)
        g3 = ROOT.TGraph(len(x3), x3, y3)
        g4 = ROOT.TGraph(len(x4), x4, y4)
        g1.SetLineColor(ROOT.kYellow)
        g2.SetLineColor(ROOT.kRed)
        g3.SetLineColor(ROOT.kBlue)
        g4.SetLineColor(ROOT.kGreen)
        mg1.Add(g1)
        mg2.Add(g2)
        mg3.Add(g3)
        mg4.Add(g4)

        if (n+1) >= options.nevents:
            break

    print

    c = ROOT.TCanvas('c', '', 800, 600)
    c.Divide(1,4)
    c.cd(1)
    mg1.Draw('AL')
    c.cd(2)
    mg2.Draw('AL')
    c.cd(3)
    mg3.Draw('AL')
    c.cd(4)
    mg4.Draw('AL')
    c.Update()

    raw_input()

if __name__ == '__main__':
    import sys
    import optparse
    usage = '%prog file'
    parser = optparse.OptionParser(usage)
    parser.add_option('-n', type='int', dest='nevents', default=1000)
    options, args = parser.parse_args()
    persist(args[0], options.nevents)
